# 若无法解析导入torch，可能是torch未安装
# 可以使用以下命令安装torch
# pip install torch torchvision
import torch
import torch.nn as nn
from torch.optim import Adam
from data_loader import train_loader, test_loader
from unet_model import UNET
import matplotlib.pyplot as plt
import numpy as np


# 初始化模型
print("begin create model...")
model = UNET(in_channels=3, out_channels=1)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=1e-4)

# 训练模型
num_epochs = 10
print("begin train...")
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, masks in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, masks)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {running_loss / len(train_loader)}')

# 测试模型
# model.eval() 会通知模型当前处于评估阶段（即测试或验证阶段），而不是训练阶段。
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for images, masks in test_loader:
        outputs = model(images)
        predicted = (torch.sigmoid(outputs) > 0.5).float()
        total += masks.numel()
        correct += (predicted == masks).sum().item()

print(f'Accuracy: {100 * correct / total}%')

# 可视化预测结果
with torch.no_grad():
    # 获取一个测试batch
    test_images, test_masks = next(iter(test_loader))
    outputs = model(test_images)
    pred_masks = (torch.sigmoid(outputs) > 0.5).float()

    # 创建可视化画布
    plt.figure(figsize=(12, 6))

    # 显示前4个样本
    for i in range(4):
        # 原始图像（反标准化处理）
        img = test_images[i].permute(1, 2, 0).numpy()
        img = img * np.array([0.229, 0.224, 0.225]) + np.array([0.485, 0.456, 0.406])  # 反标准化
        img = np.clip(img, 0, 1)

        # 真实掩码
        true_mask = test_masks[i].squeeze().numpy()

        # 预测掩码
        pred_mask = pred_masks[i].squeeze().numpy()

        # 绘制子图
        plt.subplot(3, 4, i + 1)
        plt.imshow(img)
        plt.title('Input Image')
        plt.axis('off')

        plt.subplot(3, 4, i + 5)
        plt.imshow(true_mask, cmap='gray')
        plt.title('True Mask')
        plt.axis('off')

        plt.subplot(3, 4, i + 9)
        plt.imshow(pred_mask, cmap='gray')
        plt.title('Predicted Mask')
        plt.axis('off')

    plt.tight_layout()
    plt.show()